import pygame
import sys
import traceback
import math
import time
import threading

class MandleSurface:
    # Represents a pygame surface which covers the display screen

    # EMPTY is a colour
    EMPTY=(150, 150, 150)
    def __init__(self, screen):
        self.screen=screen
        self.surface=self.init_surface()
        self.pos=( -self.screen.get_width()/2, -self.screen.get_height()/2)
        
    def init_surface(self):
        #create a grey surface to cover the display
        x,y = 2*self.screen.get_width(), 2*self.screen.get_height() 
        screen = pygame.Surface((x, y))
        background = pygame.Surface(screen.get_size())
        # create a background the same size as the window
        background = background.convert()
        # convert the background format to an appropriate one
        background.fill(self.EMPTY)
        # fill the background with grey
        screen.blit(background,(0,0))
        # blit(transfer) the background to the screen at location (0,0)
        return screen
    
    def centre(self):
        self.pos=( -self.screen.get_width()/2, -self.screen.get_height()/2)          

class MandleRender(threading.Thread):

    def __init__(self, mandle):
        threading.Thread.__init__ ( self )
        self.__mandle=mandle
        self.start()

    def run(self):
        print "Starting MandleRender thread"
        self.__mandle.running=True
        self.__mandle.draw()
        self.__mandle.stopping=False
        self.__mandle.running=False
        print "MandleRender thread ended."
    
class MandleScroll(threading.Thread):

    def __init__(self, mandle):
        threading.Thread.__init__ ( self )
        self.__mandle=mandle
        self.start()

    def run(self):
        print "Starting MandleRender thread"
        self.__mandle.running=True
        self.__mandle.re_draw()
        self.__mandle.stopping=False
        self.__mandle.running=False
        print "MandleScroll thread ended."


class Mandle():
    ZOOM_FACTOR=2
    __max_reps=0 
    reps=30
    max_reps=0
    # the origin is the centre of the display as a complex number
    origin=0+0j
    pxls_pr_unit=100
    running=True
    stopping=False
    
    def __init__(self, screen):
        self.screen=screen
        self.MSurface=MandleSurface(screen)
        

    def stop(self):
        # request stop and wait for thread to finish
        self.stopping=True
        while self.running:
             print "waiting for threads to stop"
             time.sleep(0.1)
        print "stopped"
        self.stopping=False
   
    def set_origin(self, origin):
        self.origin=origin
    
    def mandle_test(self, c):
        # tests if complex number c is part of the mandelbrot set
        # and returns the number of iterations before diverging
        # or reps if not
        z=0
        for n in xrange(self.reps):
            z=z*z+c
            if(abs(z)>4):
                return n
        return self.reps

    def scale_coords(self, coords):
        # map screen coords to a complex number
        try:
            #calculate x,y relative to origin
            x, y =float(), float()
            x, y =coords
            #convert screen coords to surface coords
            width=self.screen.get_width()
            height=self.screen.get_height()
            z=complex(x-width/2, y-height/2)/self.pxls_pr_unit
            z=z + self.origin         
        except TypeError, e:
            print "pxls_pr_unit, coords, origin\n"
            print self.pxls_pr_unit, self.coords, self.origin
            print "\n"
            raise
            
        return z

    def get_colour(self, n, reps ):
        # derive an (R, G, B) colour tuple from a number n
        if n==reps:
            #BLACK
            return (0, 0, 0)
        n=8*abs(n % 32)
        if n < 64 :
            #red->yellow
            red=255
            green=4*n
            blue=0
        elif n < 128 :
            #yellow->green
            red=4*(128-n)-1
            green=255
            blue=0
        elif n < 192 :
            #green->blue
            red=0
            green=4*(192-n)-1
            blue=4*(n-128)
        else:
           #blue->red
           red=4*(n-192)
           #green=4*(256-n)
           green=0
           blue=4*(256-n)-1
        if red >255 or green > 255 or blue > 255:
           raise Exception( "Invalid Colour (R,G,B)" , (red, green, blue) )
        return red, green, blue
    
    def same(self, sx, sy, step):
        # calculates if neighbouring squares
        # are the same colour
        try:
            col=self.MSurface.surface.get_at((sx, sy-step))
            same=col == \
            self.MSurface.surface.get_at((sx, sy+step))== \
            self.MSurface.surface.get_at((sx+step, sy))== \
            self.MSurface.surface.get_at((sx-step, sy))
            if same : return col
        except IndexError, e:
            pass
        return None
                
    def draw(self):
        # render a mandelbrot set on a surface
        # paint the surface on the screen
        print "draw(self):"
        print "origin:", self.origin
        self.MSurface=MandleSurface(self.screen)
        win_width=self.screen.get_width()
        win_height=self.screen.get_height()
        for step in (8,4,2,1):
            x_axis=xrange(0, win_width , step )         
            y_axis=xrange(0, win_height, step )
            # make sure the origin is in the middle of the screen
            self.MSurface.centre() 
            self.__max_reps=0
            #pygame.draw.rect(Surface, color, Rect, width=0): return Rect
            skip=False
            for x in x_axis:
                if step < 8 : skip=True
                for y in y_axis:
                    if self.stopping==True :
                        print "Thread stopping"
                        return
                    sx, sy= x+win_width/2, y+win_height/2
                    col=self.same(sx, sy, step)
                    if step == 8 or col==None :
                        z = self.scale_coords( (x, y) )
                        result=self.mandle_test(z)
                        col=self.get_colour(result, self.reps)  
                        if(result < self.reps):
                            self.__max_reps=max((self.__max_reps, result))
                    self.MSurface.surface.fill( col, \
                    pygame.rect.Rect( sx, sy, step , step ))           
                # update the display
                self.screen.blit(self.MSurface.surface, self.MSurface.pos)
                pygame.display.flip()
        
    def re_draw(self):
            # loop over the screen and fill in any grey bits
            print "starting re_draw"       
            newSurface=MandleSurface(self.screen)
            newSurface.centre()
            win_width, win_height=self.screen.get_width(),self.screen.get_height()
            x_axis, y_axis=xrange(win_width), xrange(win_height)
            for x in x_axis:
                for y in y_axis:
                    if self.stopping==True :
                        print "Thread stopped"
                        return
                    sx, sy = x-newSurface.pos[0], y-newSurface.pos[1]
                    curr_col=self.screen.get_at((x, y))
                    if curr_col==MandleSurface.EMPTY:
                        z = self.scale_coords( (x, y) )
                        result=self.mandle_test(z)                    
                        col=self.get_colour(result, self.reps)                 
                        self.screen.set_at( (x, y) , col )
                        newSurface.surface.set_at( (sx, sy) , col )
                    else:                     
                        newSurface.surface.set_at( (sx, sy) , curr_col )
                # update the display               
                pygame.display.flip()
            #swap the old surface for a new one
            self.screen.blit(newSurface.surface, newSurface.pos)
            pygame.display.flip()
            self.MSurface=newSurface         
        
    def zoom(self):
        self.pxls_pr_unit=int(self.pxls_pr_unit*self.ZOOM_FACTOR) 
        # adjust reps up or down
        if self.__max_reps+1==self.reps:
            self.reps=2*self.reps
        else:
            self.reps=int((self.__max_reps+self.reps)/2)
        # update surface_pos
        
        print "reps:", self.reps, "pxls_pr_unit:", self.pxls_pr_unit
        print "print mandle finished!"
        print "origin:", self.origin, "\n"
        
    def scroll(self, dx, dy):
        # move the surface around the screen
        print "scrolling:", dx, dy

        #Don't let the surface scroll off the screen
        screen_rect=self.screen.get_rect()
        display_rect=self.MSurface.surface.get_rect()
        display_rect.move(dx, dy)
        if not display_rect.colliderect(screen_rect):
            return
       
        # repaint
        dest=(self.MSurface.pos[0]+dx, self.MSurface.pos[1]+dy)
        self.screen.blit(self.MSurface.surface, dest)
        pygame.display.flip()

        # update variables
        self.MSurface.pos=dest
        self.origin-=complex(float(dx), float(dy))/self.pxls_pr_unit

        
def clear_screen(screen):
    background = pygame.Surface(screen.get_size())
    # create a background the same size as the window
    background = background.convert()
    # convert the background format to an appropriate one
    background.fill((0,0,0))
    # fill the background with black
    screen.blit(background,(0,0)) # blit(transfer) the background to the screen at location (0,0)
    pygame.display.flip()
    
def main():
    # parameters
    win_width = 800
    win_height = 400
    screen = pygame.display.set_mode((win_width, win_height))
    pygame.display.set_caption("Drag to explore, double-click to zoom")
    running = True
    clock=pygame.time.Clock ()
    
    clear_screen(screen)
    my_mandle=Mandle(screen)

    # draw the whole set
    MandleRender(my_mandle)

    # Used for detecting double-clicks.
    #            (time,            (x, y), button)
    last_click = (time.clock(),    (0, 0), -1    )
    drag=False       
    while running: 
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                 running = False
            elif event.type == pygame.MOUSEMOTION:
                coords=event.dict['pos']
                #print "cursor->", my_mandle.scale_coords((coords))
                #print event
                dx, dy=event.dict['rel']
                if event.dict['buttons'][0]==1:
                    my_mandle.stop()
                    my_mandle.scroll(dx, dy)
                    pygame.display.flip()
                    drag=True
            elif event.type == pygame.MOUSEBUTTONUP:
                coords=event.dict['pos']
                #print my_mandle.scale_coords((coords))
                now=time.clock()
                if drag:
                    #dragging must have finished
                    #colour in any grey bits
                    my_mandle.stop()
                    MandleScroll(my_mandle)
                    drag=False
                elif now -last_click[0] < 0.5:
                    # Double click!
                    # calculate the new origin to center the image
                    my_mandle.stop()
                    origin=my_mandle.scale_coords((coords))
                    # move the origin 
                    my_mandle.set_origin(origin)
                    # zoom in
                    my_mandle.zoom()                 
 
                    # re-draw
                    #stop drawing thread if one is running

                    #clear_screen(screen)
                    #create a new thread to render mandle
                    MandleRender(my_mandle)
                last_click = ( now, (0, 0), 0)
        
        #refresh rate
        clock.tick(30)


try:
    main()
except Exception, e:
    tb = sys.exc_info()[2]
    traceback.print_exception(e.__class__, e, tb)
pygame.quit()
